from langchain_core.messages import HumanMessage
from langgraph.graph import StateGraph

from owl_ai.domain.agent_config_entity import AgentConfigEntity
from owl_ai.service.agent_flow_service import State, chat_llm_generate
from owl_ai.service.graph.graph_compile import GraphCompile, ChatNode


class AppWorkFlowCompile(GraphCompile):

    @classmethod
    def compile(cls, config: AgentConfigEntity):
        """
        应用类工作流编排
        Args:
            config:

        Returns:

        """
        graph = StateGraph(State)

        agent_config = config.config

        for node_name, node_config in agent_config.items():
            node_llm_config = node_config.get("llm")
            if node_llm_config:
                chat_llm = GraphCompile.chat_llm_generate(node_llm_config, tools=None)
                graph_chat_node = ChatNode(chat_llm=chat_llm, system_prompt=node_config.get("systemPrompt"))
                graph.add_node(node_name, graph_chat_node)

        for node_name, node_config in agent_config.items():
            edges_config = node_config.get("edges")
            out_edge_type = edges_config.get("type")

            if "normal" == out_edge_type:
                out_node_name = edges_config.get("outNode")
                graph.add_edge(node_name, out_node_name)

        return graph.compile()
